import gradio as gr
import os
import base64
import requests
import io
from PIL import Image

# 定义上传接口的 URL
upload_url = "http://106.54.192.233:8080/uploadAsset"

# 定义获取图片列表接口的 URL
get_image_url = "http://106.54.192.233:8080/getImageList"

# 定义请求头
headers = {
    "Content-Type": "application/json"
}

def upload_image(image, name, overwrite):
    if image is None:
        return "请上传图片文件"
    if not name:
        return "请指定文件的名称"
    
    # 将 numpy.ndarray 转换为文件对象并进行 Base64 编码
    try:
        pil_image = Image.fromarray(image)  # 将 numpy.ndarray 转为 PIL.Image
        buffer = io.BytesIO()              # 创建字节流缓冲区
        pil_image.save(buffer, format="PNG")  # 保存图片到缓冲区
        encoded_file = base64.b64encode(buffer.getvalue()).decode('utf-8')  # Base64 编码
    except Exception as e:
        return f"图片编码失败: {e}"
    
    # 构建请求体
    data = {
        "name": name,
        "type": "image/png",                 # 文件类型为 PNG
        "asset": f"data:image/png;base64,{encoded_file}",  # 带格式的 Base64 编码
        "overwrite": overwrite              # 是否覆盖
    }
    
    # 发送 POST 请求
    response = requests.post(upload_url, json=data, headers=headers)

    # 检查响应状态码
    if response.status_code == 200:
        response_data = response.json()
        return f"文件 {name} 上传成功: {response_data}"
    else:
        return f"文件 {name} 上传失败，状态码: {response.status_code}, 响应: {response.text}"

def download_image(name):
    if not name:
        return None, "请指定文件名称"
    
    # 构建请求体
    data = {"names": [name]}
    
    # 发送 POST 请求
    response = requests.post(get_image_url, json=data, headers=headers)

    # 检查响应状态码
    if response.status_code == 200:
        response_data = response.json()
        if name in response_data:
            image_url = response_data[name]
            image_url = "http://106.54.192.233/assets/" + image_url
            try:
                # 下载图片
                image_response = requests.get(image_url)
                if image_response.status_code == 200:
                    image = Image.open(io.BytesIO(image_response.content))
                    return image, f"文件 {name} 下载成功"
                else:
                    return None, f"无法下载图片，状态码: {image_response.status_code} + {image_url}"
            except Exception as e:
                return None, f"图片下载失败: {e}"
        else:
            return None, "指定的图片名称未找到"
    else:
        return None, f"获取图片列表失败，状态码: {response.status_code}, 响应: {response.text}"

# 定义 Gradio 界面
def reset_overwrite():
    return False

with gr.Blocks() as demo:
    with gr.Tab("上传"):
        gr.Markdown("### 上传图片")
        with gr.Row():
            image_input = gr.Image(label="上传图片", type="numpy")  # 返回 numpy 数组
        with gr.Row():
            name_input = gr.Textbox(label="指定文件名", placeholder="请输入文件名")
        with gr.Row():
            overwrite_checkbox = gr.Checkbox(label="覆盖已存在的文件", value=False)
        with gr.Row():
            upload_button = gr.Button("上传")
        result_output = gr.Textbox(label="上传结果", interactive=False)
        
        # 绑定上传功能
        upload_button.click(
            fn=upload_image,
            inputs=[image_input, name_input, overwrite_checkbox],
            outputs=result_output
        )
        
        # 上传后重置 overwrite_checkbox
        upload_button.click(
            fn=reset_overwrite,
            inputs=[],
            outputs=[overwrite_checkbox]
        )

    with gr.Tab("下载"):
        gr.Markdown("### 下载图片")
        with gr.Row():
            download_name_input = gr.Textbox(label="图片名称", placeholder="请输入图片名称")
        with gr.Row():
            download_button = gr.Button("下载")
        with gr.Row():
            downloaded_image = gr.Image(label="下载的图片", interactive=False)
        with gr.Row():
            download_log = gr.Textbox(label="下载日志", interactive=False)
        
        # 绑定下载功能
        download_button.click(
            fn=download_image,
            inputs=[download_name_input],
            outputs=[downloaded_image, download_log]
        )

# 启动 Gradio 应用
demo.launch()
